﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Threading;

namespace Runtime.Common.Utility
{
    public static class TaskPool
    {
        internal static int MaxPoolSize;

        // avoid to use ConcurrentDictionary for safety of WebGL build.
        private static readonly ConcurrentDictionary<Type, Func<int>> Sizes =
            new ConcurrentDictionary<Type, Func<int>>();

        static TaskPool()
        {
            try
            {
                var value = Environment.GetEnvironmentVariable("UNITASK_MAX_POOLSIZE");
                if (value != null)
                {
                    if (int.TryParse(value, out var size))
                    {
                        MaxPoolSize = size;
                        return;
                    }
                }
            }
            catch
            {
            }

            MaxPoolSize = int.MaxValue;
        }

        public static void SetMaxPoolSize(int maxPoolSize)
        {
            MaxPoolSize = maxPoolSize;
        }

        public static IEnumerable<(Type, int)> GetCacheSizeInfo()
        {
            foreach (var item in Sizes)
            {
                yield return (item.Key, item.Value());
            }
        }

        public static void RegisterSizeGetter(Type type, Func<int> getSize)
        {
            Sizes[type] = getSize;
        }
    }

    public interface ITaskPoolNode<T>
    {
        ref T NextNode { get; }
    }

    [StructLayout(LayoutKind.Auto)]
    public struct TaskPool<T>
        where T : class, ITaskPoolNode<T>
    {
        int gate;
        int size;
        T root;

        public int Size => size;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool TryPop(out T result)
        {
            if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
            {
                var v = root;
                if (!(v is null))
                {
                    ref var nextNode = ref v.NextNode;
                    root = nextNode;
                    nextNode = null;
                    size--;
                    result = v;
                    Volatile.Write(ref gate, 0);
                    return true;
                }

                Volatile.Write(ref gate, 0);
            }

            result = default;
            return false;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool TryPush(T item)
        {
            if (Interlocked.CompareExchange(ref gate, 1, 0) == 0)
            {
                if (size < TaskPool.MaxPoolSize)
                {
                    item.NextNode = root;
                    root = item;
                    size++;
                    Volatile.Write(ref gate, 0);
                    return true;
                }
                else
                {
                    Volatile.Write(ref gate, 0);
                }
            }

            return false;
        }
    }

    internal sealed class PooledDelegate<T> : ITaskPoolNode<PooledDelegate<T>>
    {
        static TaskPool<PooledDelegate<T>> pool;

        PooledDelegate<T> nextNode;
        public ref PooledDelegate<T> NextNode => ref nextNode;

        static PooledDelegate()
        {
            TaskPool.RegisterSizeGetter(typeof(PooledDelegate<T>), () => pool.Size);
        }

        readonly Action<T> runDelegate;
        Action continuation;

        PooledDelegate()
        {
            runDelegate = Run;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public static Action<T> Create(Action continuation)
        {
            if (!pool.TryPop(out var item))
            {
                item = new PooledDelegate<T>();
            }

            item.continuation = continuation;
            return item.runDelegate;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        void Run(T _)
        {
            var call = continuation;
            continuation = null;
            if (call != null)
            {
                pool.TryPush(this);
                call.Invoke();
            }
        }
    }
}